import torch.nn as nn
from typing import Tuple

from transformers import BertModel, BertForSequenceClassification


class BertClassifier(nn.Module):
    def __init__(self,
                 model_dir: str,
                 num_labels: int,
                 dropout_prob=0.1):
        super(BertClassifier, self).__init__()
        self.bert = BertModel.from_pretrained(model_dir)
        self.linear = nn.Linear(self.bert.config.hidden_size, num_labels)
        self.dropout = nn.Dropout(dropout_prob)

    def forward(self,
                inputs: Tuple,
                # task_name: str = 'main',  # auxiliary
                ):
        outputs = self.bert(input_ids=inputs[0], attention_mask=inputs[1])
        pooled_output = self.dropout(outputs[1])
        logits = self.linear(pooled_output)
        return logits, pooled_output
        # if task_name == 'main':
        #     logits = self.linear(pooled_output)
        #     return logits
        # elif task_name == 'auxiliary':
        #     return pooled_output
